from elasticsearch_dsl import connections
from tqdm import tqdm


def insert_into_es(chunks, index_name="insurance_manual_es_01"):
    """
    将文本片段插入到ES索引

    参数:
        chunks: 文本片段列表（包含page_num和chunk）
        index_name: 索引名称
    """
    # 连接ES
    connections.create_connection(hosts=["http://localhost:9200"])

    # 从ES获取文档类（需要和初始化时的结构一致）
    from elasticsearch_dsl import Document, Text, Integer
    class InsuranceDocument(Document):
        page_num = Integer()
        text = Text()

        class Index:
            name = index_name

    # 批量插入数据
    for chunk in tqdm(chunks, desc="插入ES中"):
        doc = InsuranceDocument(
            page_num=chunk["page_num"],
            text=chunk["chunk"]
        )
        doc.save()  # 保存到ES

    # 刷新索引（确保数据可查）
    InsuranceDocument._index.refresh()
    print(f"成功向ES插入 {len(chunks)} 条数据")


# 测试代码
if __name__ == "__main__":
    # 模拟文本片段
    sample_chunks = [
        {"page_num": 1, "chunk": "平安保险成立于1988年，总部位于深圳。"},
        {"page_num": 1, "chunk": "旗下有多个保险产品，包括寿险、车险、健康险等。"},
        {"page_num": 2, "chunk": "平安福是公司的明星寿险产品，2024年推出了新版。"}
    ]

    # 插入ES
    insert_into_es(sample_chunks)
